from __future__ import division
from data.bbox.bbox_dataset import DetectionDataset
import os
import numpy as np
import json
class txtDetDataset(DetectionDataset):
    def __init__(self,anno_path, img_root, threshold = .7, classes = ("person",)):
        super(txtDetDataset,self).__init__()
        a = open(anno_path,"rt")
        objs = {}
        for line in a:
            try:
                name,x0,y0,x1,y1,score,cls = line.strip().split("\t")
            except:
                continue
            if float(score) < threshold :continue
            x0 = float(x0)
            y0 = float(y0)
            x1 = float(x1)
            y1 = float(y1)
            cls = int(cls) - 1
            try:
                objs[name].append([x0,y0,x1,y1,cls,0])
            except KeyError:
                objs[name] = [[x0, y0, x1, y1, cls, 0]]

        self.objs = [(k,objs[k]) for k in objs.keys()]
        self.classes  = ["bike","bus","car","motor","person","rider","traffic light","traffic sign","train","truck"]
        self.img_root = img_root
    def at_with_image_path(self, idx):
        name = self.objs[idx][0]
        bboxes = self.objs[idx][1]
        print(bboxes)
        return os.path.join(self.img_root,name), np.array(bboxes).astype(float)
    def __len__(self):
        return len(self.objs)
if __name__ == '__main__':
    # txtDetDataset("/data1/zyx/yks/sources/Detectron/output/results_1538288326.39.txt",
    #               "/data1/zyx/yks/dataset/retail/test_images/").viz()
    txtDetDataset("/data1/zyx/yks/sources/Detectron/output/retail_retina/results_1538445336.05.txt",
                  "/data1/zyx/yks/dataset/retail/test_images", threshold=.5).viz()